/** \file    refine/smooth_poisson.cpp
  * \brief   Construct smooth grey mask from existing bit mask solving Poisson equation.
  * \date    2005
  * \author  Fedor Chelnokov
  * \version $Header: /home/cscom1r4/cvsroot/cs/refine/smooth_poisson.cpp,v 1.2 2005/06/25 18:49:42 cscom1r4 Exp $
  */

#include <limits>
#include <iostream>
#include "refine/smooth_poisson.h"
#include "imglib/imageclass.h"
#include "imglib/bitmap.h"
#include "imglib/stat_calculator.h"
#include "imglib/dichromatic_raster.h"
#include "imglib/bmpmorph.h"
#include "refine/topology_limit.h"
#include "xtd/minmax"
#include <xtd/tree2d.h>

using std::clog;
using std::endl;

namespace refine {

namespace smooth_poisson {

const params default_params = {5, 1.0f, false, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;

template <class R, class M>
class smoother {
public:
	typedef typename R::elem_type color_type;
public:
	const params & _params;
public:
	smoother(const params & p, const R & s, const bit_mask::mask & m) 
		: _params(p), _src(s), _src_mask(&m) 
		{ }
	void run(M & b);

	const R & src() const
		{ return _src; }
	const color_type & src(size_point p) const
		{ return _src[p]; }
	const color_type & src(int_point p) const
		{ return _src[p]; }
	const color_type & src(unsigned x, unsigned y) const
		{ return _src(x,y); }
	unsigned width() const
		{ return _src.width(); }
	unsigned height() const
		{ return _src.height(); }
	size_extent dim() const
		{ return _src.dim(); }

	const bit_mask::mask & src_mask() const
		{ return *_src_mask; }
	bool src_mask(size_point p) const
		{ return (*_src_mask)[p]; }
	bool src_mask(int_point p) const
		{ return (*_src_mask)[p]; }
	bool src_mask(unsigned x, unsigned y) const
		{ return (*_src_mask)(x,y); }

	const M & smoothed() const
		{ return *_smoothed; }
	grey_mask_value smoothed(size_point p) const
		{ return (*_smoothed)[p]; }
	grey_mask_value smoothed(int_point p) const
		{ return (*_smoothed)[p]; }
	grey_mask_value smoothed(unsigned x, unsigned y) const
		{ return (*_smoothed)(x,y); }
private:
	const R & _src;
	const bit_mask::mask * _src_mask;
	M * _smoothed;
	int_rect _frame;
};

template <class R, class M>
static void perform_impl(const R & src, const bit_mask::mask & m, M & s, const params & p)
{
	if (!p.has_initial_apprx)
		s = img::make_dichromatic_raster(m, grey_mask_min_value, grey_mask_max_value);
	smoother<R, M>(p, src, m).run(s);
}

void perform(const byte_image & src, const bit_mask::mask & m, grey_mask::mask & s, const params & p)
	{ perform_impl(src, m, s, p); }

void perform(const argb_byte_image & src, const bit_mask::mask & m, grey_mask::mask & s, const params & p)
	{ perform_impl(src, m, s, p); }

void perform(byte_image & img, const bit_mask::mask & m, const params & p)
	{ perform_impl(img, m, img::projection(img, &byte_color::a), p); }

void perform(argb_byte_image & img, const bit_mask::mask & m, const params & p)
	{ perform_impl(img, m, img::projection(img, &argb_byte_color::a), p); }

// smoother implementation //

struct boundary {
	int_point p; ///< point on the left or on the ceil of the boundary
	img::BGR<float> dc; ///< difference between mean colors at each side of the boundary
	operator int_point const & () const
		{ return p; }
};

template <class R, class M>
void smoother<R, M>::run(M & b)
{
	_smoothed = & b;
	_frame = int_rect(0, 0, width(), height());

	assert(src_mask().dim() == dim());
	assert(_smoothed->dim() == dim());

	img::bitmap fixed = *_src_mask; // alpha for fixed points will not change
	img::erode8(fixed, _params.window_half_size);
	{
		img::bitmap tmp = *_src_mask;
		tmp.inverse();
		img::erode8(tmp, _params.window_half_size);
		img::combine(fixed, tmp, img::bit_or());
	}

	typedef xtd::tree2d<boundary, int_point, int_rect> tree_type;
	tree_type boundary_tree; ///< contains calculated statistics for every boundary point

	//building of boundary tree
	for (img::bitmap::const_iterator it = src_mask().begin(); it != src_mask().end(); ++it)
	{
		size_point p = it;
		if (p.x + 1 == width() || p.y + 1 == height())
			continue;
		if (*it == it.rel(1,0) && *it == it.rel(0,1))
			continue;
		boundary b;
		b.p = p;

		img::stat_calculator calc[2];
		int_point d;
		for (d.y = -_params.window_half_size; d.y <= _params.window_half_size; ++d.y) {
			for (d.x = -_params.window_half_size; d.x <= _params.window_half_size; ++d.x) {
				int_point z = int_point(p) + d;
				if (!_frame.inside(z))
					continue;
				int i = refine::ii(src_mask(z));
				calc[i].add(src(z));
			}
		}
		assert (calc[0].count() > 0 && calc[1].count() > 0);
		b.dc = calc[1].mean() - calc[0].mean();
		boundary_tree.push_back(b);
	}
	boundary_tree.rebalance();

	typedef img::image<short> force_image;
	short const force_min = std::numeric_limits<short>::min();
	short const force_max = std::numeric_limits<short>::max();
	force_image force(dim()); //we solve equation at each pixel: \laplacian \alpha = force

	unsigned const max_distance = _params.window_half_size * _params.window_half_size;
	unsigned roi_area = 0;
	for (img::bitmap::iterator it = fixed.begin(); it != fixed.end(); ++it) {
		if (*it)
			continue;
		size_point p = it;
		if (p.x == 0 || p.y == 0 || p.x + 1 == width() || p.y + 1 == height()) {
			*it = true;
			continue; //boundary points are always fixed
		}
		++roi_area;

		//find nearest point in boundary tree
		int_rect r(p, p);
		r.inflate(_params.window_half_size+1, _params.window_half_size+1);
		unsigned min_distance = std::numeric_limits<unsigned>::max();
		img::BGR<float> dc;
		for (tree_type::range_iterator ri = boundary_tree.range_begin(r); 
			ri != boundary_tree.range_end(); ++ri) 
		{
			boundary const & b = *ri;
			unsigned dist = 
				(int(p.x) - int(b.p.x))*(int(p.x) - int(b.p.x)) +
				(int(p.y) - int(b.p.y))*(int(p.y) - int(b.p.y));
			if (dist < min_distance) {
				min_distance = dist;
				dc = b.dc;
			}
		}
		assert (min_distance < std::numeric_limits<unsigned>::max()); //range was not empty

		float dd = dc.sqr<float>();

		if (dd == 0 || min_distance >= max_distance)
			force[it] = 0;
		else {
			img::BGR<float> e = src(it);
			e *= -4;
			e += src(p.x - 1, p.y);
			e += src(p.x + 1, p.y);
			e += src(p.x, p.y - 1);
			e += src(p.x, p.y + 1);

			float f = grey_mask_max_value * _params.adaptability * //(1 - sqrt(float(min_distance) / max_distance)) *
				(dc.b * e.b + dc.g * e.g + dc.r * e.r) / dd;
			if (f < force_min)
				force[it] = force_min;
			else if (f > force_max)
				force[it] = force_max;
			else
				force[it] = f;
		}
	}

	if (!_params.silent) {
		clog << "Poisson mask smoothing initiated...\n"
			"region of interest area = " << roi_area << endl;
	}

	img::bitmap curr(dim()), next(dim());
	curr.fill(true);

	for (int pass = 0; pass < 250; ++pass, curr.swap(next)) {
		unsigned tests = 0, changes = 0;
		next.fill(false);

		for (img::bitmap::const_iterator it = fixed.begin(); it != fixed.end(); ++it) {
			if (*it || !curr[it])
				continue;
			++tests;
			size_point p = it;

			grey_mask_value min_nei, max_nei;
			min_nei = grey_mask_min_value;
			max_nei = grey_mask_max_value;
			//refine::grey_mask::topology_limit(smoothed(), p, min_nei, max_nei);
			if (min_nei > max_nei)
				continue;

			int v = smoothed(p.x - 1, p.y);
			v += smoothed(p.x + 1, p.y);
			v += smoothed(p.x, p.y - 1);
			v += smoothed(p.x, p.y + 1);
			v -= force[p];
			v /= 4;
			if (v < min_nei)
				v = min_nei;
			else if (v > max_nei)
				v = max_nei;
			if (v != (*_smoothed)[p]) {
				++changes;
				(*_smoothed)[p] = v;
				next[p - size_point(-1,1)] = true;
				next[p - size_point(0,1)] = true;
				next[p - size_point(1,1)] = true;
				next[p - size_point(1,0)] = true;
				curr[p + size_point(-1,1)] = true;
				curr[p + size_point(0,1)] = true;
				curr[p + size_point(1,1)] = true;
				curr[p + size_point(1,0)] = true;
			}
		}

		if (!_params.silent)
			clog << "pass " << pass <<
					": tests=" << tests << 
					", changes=" << changes << 
					endl;

		if (changes == 0)
			break;
	}

	if (!_params.silent)
		clog << endl;
}

} //namespace smooth_poisson

} //namespace refine
